"""访问计算模块 - 用于计算卫星和地面目标之间的访问窗口"""

import numpy as np
from astropy import units as u
from astropy.time import Time, TimeDelta
from poliastro.bodies import Earth
from poliastro.core.spheroid_location import cartesian_to_ellipsoidal
from poliastro.frames import Planes
from typing import Union, List, TYPE_CHECKING, Any, Tuple

import concurrent.futures
from functools import partial

# 使用TYPE_CHECKING来处理类型注解，避免循环导入
if TYPE_CHECKING:
    from satellite import Satellite
    from target import GroundTarget
else:
    # 运行时不需要实际的类，只需要有名称就行
    Satellite = Any
    GroundTarget = Any

from abc import ABC, abstractmethod
# 添加Observed抽象基类  
class Observed(ABC):
    """被观察者抽象基类，定义了被观察者（卫星或地面目标）的通用接口"""
    
    @abstractmethod
    def get_itrs_xyz(self, time:Time = None) -> u.Quantity:
        """
        被观测对象的itrs坐标
        
        参数:
            time: 时间点
        返回:
            itrs坐标
        """
        pass
# 添加Observer抽象基类  
class Observer(ABC):
    """观察者抽象基类，定义了观察者（卫星或地面目标）的通用接口"""
    
    @abstractmethod
    def can_see(self, target:Observed, time:Time = None) -> Tuple[bool]:
        """
        判断观察者是否可以看到目标
        
        参数:
            target: 观测目标(卫星或地面站)的itrs_xyz坐标
            time: 时间点
        返回:
            Tuple[bool]: (是否可见)
        """
        pass


class AccessWindow:
    """访问窗口类"""
    
    def __init__(self, observer, observed, start_time, end_time, trajectory_points=None):
        """
        初始化访问窗口
        
        参数:
            observer: 观察者（卫星或地面目标）
            target: 观测目标（地面目标或卫星）
            start_time: 窗口开始时间
            end_time: 窗口结束时间，可以为None表示窗口尚未结束
            trajectory_points: 轨迹点字典，包含时间、经纬度和可见性信息
        """
        self.observer = observer
        self.observed = observed
        self.start_time = start_time
        self.end_time = end_time
        
        # 处理end_time为None的情况
        if end_time is not None:
            self.duration = end_time - start_time
        else:
            # 如果end_time未知，duration也暂时未知
            self.duration = None
            
        self.trajectory_points = trajectory_points or {}
        
    def __repr__(self):
        """字符串表示"""
        observer_name = getattr(self.observer, 'name', 'Unknown')
        observed_name = getattr(self.observed, 'name', 'Unknown')
        
        # 处理 end_time 和 duration 可能为 None 的情况
        end_time_str = f"end={self.end_time.isot}, " if self.end_time is not None else "end=未结束, "
        duration_str = f"duration={self.duration})" if self.duration is not None else "duration=未确定)"
        
        return (f"AccessWindow(observer='{observer_name}', "
                f"observed='{observed_name}', "
                f"start={self.start_time.isot}, "
                f"{end_time_str}"
                f"{duration_str}")
                
                
from astropy.coordinates import GCRS, ITRS, CartesianRepresentation
class AccessCalculator:
    """访问窗口计算器"""

    @staticmethod
    def transform_to_itrs(satellite_orbit)->ITRS:
        """
        将卫星轨道换到ITRS坐标系
        return ITRS
        """

        # 确保r_sat和r_target都在ITRS坐标系中
        if hasattr(satellite_orbit, 'plane'):
            # 如果卫星轨道是GCRS坐标系
            if satellite_orbit.plane == Planes.EARTH_EQUATOR:
                # 创建GCRS坐标对象
                gcrs = GCRS(
                    satellite_orbit.r,
                    obstime=satellite_orbit.epoch,
                    representation_type=CartesianRepresentation,
                )
                # 将GCRS坐标转换为ITRS坐标
                itrs = gcrs.transform_to(ITRS(obstime=satellite_orbit.epoch))

        else:        
            raise ValueError(f"Unsupported plane!")
        return itrs

    @staticmethod
    def calculate_elevation(target_itrs_xyz:u.Quantity, ground_itrs_xyz:u.Quantity):
        """
        计算卫星相对于地面目标的仰角
        
        参数:
            target_itrs_xyz: 目标(卫星或地面目标)(ITRS坐标系中的xyz坐标)
            ground_itrs_xyz: 地面目标对象(ITRS坐标系中的xyz坐标)
            
        返回:
            仰角值 (角度)
        """
        # 注意：轨道递推应该在调用此方法前完成，避免重复递推
        # 获取卫星位置
        r_sat = target_itrs_xyz
        
        # 获取目标位置 (ITRS坐标系中)
        r_ground_target = ground_itrs_xyz
        
        # 计算目标到卫星的向量
        r_target_to_sat = r_sat - r_ground_target
        
        # 计算仰角
        # 仰角是地平面与卫星连线之间的夹角
        # 首先需要确定地平面的法向量（即目标位置的单位向量）
        r_target_unit = r_ground_target / np.linalg.norm(r_ground_target.value)
        
        # 目标到卫星的单位向量
        r_target_to_sat_unit = r_target_to_sat / np.linalg.norm(r_target_to_sat.value)
        
        # 计算角度（与天顶方向的夹角）
        cos_zenith = r_target_unit.dot(r_target_to_sat_unit).value
        # 确保cos_zenith在[-1, 1]范围内，避免数值误差导致的问题
        cos_zenith = max(min(cos_zenith, 1.0), -1.0)
        zenith_angle = np.arccos(cos_zenith)
        
        # 仰角 = 90度 - 与天顶方向的夹角
        elevation = 90 * u.deg - zenith_angle * u.rad
        
        return elevation
    
    @staticmethod
    def calculate_angle(satellite_itrs_xyz, ground_itrs_xyz):
        """
        计算目标相对于卫星的角度
        
        参数:
            satellite_itrs_xyz: 卫星(ITRS坐标系中的xyz坐标)
            ground_itrs_xyz: 地面目标(ITRS坐标系中的xyz坐标)
            
        返回:
        计算地面目标相对于卫星的角度
        """
        r_sat = satellite_itrs_xyz
        r_target = ground_itrs_xyz
        r_sat_to_target = r_target - r_sat 
        # 单位向量
        r_sat_to_target_unit = r_sat_to_target / np.linalg.norm(r_sat_to_target.value)
        r_sat_unit = -r_sat / np.linalg.norm(r_sat.value)
        # 计算角度
        angle = np.arccos(r_sat_to_target_unit.dot(r_sat_unit).value) * u.rad

        return angle.to(u.deg)

    @staticmethod
    def is_visible_on_the_ground(target_itrs_xyz, ground_itrs_target, min_elevation=5.0):
        """
        判断卫星在指定时间是否可见目标
        
        参数:
            target_itrs_xyz: 目标(卫星或地面目标)(ITRS坐标系中的xyz坐标)
            ground_itrs_target: 地面目标对象(ITRS坐标系中的xyz坐标) 
            min_elevation: 最小仰角(度)，默认为5度
            
        返回:
            布尔值，表示是否可见
        """
        # 计算仰角
        elevation = AccessCalculator.calculate_elevation(target_itrs_xyz, ground_itrs_target)
        
        # 检查仰角是否达到最小要求
        # 注意：此方法依赖于calculate_elevation已经考虑了地球遮挡
        return elevation >= min_elevation * u.deg
    
    @staticmethod
    def calculate_access_windows(
        observer: Observer,
        observed: Observed, 
        start_time: Time,
        end_time: Time,
        step: TimeDelta,
        min_elevation=5.0
    ) -> List[AccessWindow]:
        """
        计算可见窗口
        
        Args:
            observer: 观测者(卫星或地面站)
            observed: 观测目标(卫星或地面站) 
            start_time: 开始时间
            end_time: 结束时间
            step: 时间步长
            min_elevation: 最小仰角(度)，默认为5度
            precomputed_data: 预计算的轨道数据，可以是卫星或目标的轨道数据
            
        Returns:
            可见窗口列表
        """
        windows = []
        current_window = None
        
        # 存储所有轨迹点信息，用于可视化
        all_trajectory_points = {
            'times': [],
            'lats': [],
            'lons': [],
            'visible': []
        }
        
        # 当前窗口的轨迹点
        current_window_points = None
        
        # 遍历时间步长
        current_time = start_time
        from astropy.coordinates import SphericalRepresentation

        while current_time <= end_time:
                
            is_visible = observer.can_see(observed, current_time)
            
            # 存储轨迹点信息（如果观测者是卫星）
            if hasattr(observer, 'orbit'):
                # 从can_see返回的r_sat中获取卫星坐标（已经在ITRS坐标系中）
                # 提取经纬度
                itrs = observer.find_itrs_from_precomputed_data(current_time)
                lat = itrs.represent_as(SphericalRepresentation).lat.to(u.deg).value
                lon = itrs.represent_as(SphericalRepresentation).lon.to(u.deg).value

                # 添加到全局轨迹点
                all_trajectory_points['times'].append(current_time)
                all_trajectory_points['lats'].append(lat)
                all_trajectory_points['lons'].append(lon)
                all_trajectory_points['visible'].append(is_visible)
                
                # 如果正在一个可见窗口中，添加到当前窗口的轨迹点
                if current_window_points is not None:
                    current_window_points['times'].append(current_time)
                    current_window_points['lats'].append(lat)
                    current_window_points['lons'].append(lon)
                    current_window_points['visible'].append(is_visible)
            
            # 处理访问窗口
            if is_visible and current_window is None:
                # 开始新窗口
                current_window = AccessWindow(
                    observer=observer,
                    observed=observed,
                    start_time=current_time,
                    end_time=None  # 初始化时不确定结束时间
                )
                
                # 为新窗口创建轨迹点字典
                if hasattr(observer, 'orbit'):
                    current_window_points = {
                        'times': [],
                        'lats': [],
                        'lons': [],
                        'visible': []
                    }
                    
                    # 添加当前点
                    current_window_points['times'].append(current_time)
                    current_window_points['lats'].append(lat)
                    current_window_points['lons'].append(lon)
                    current_window_points['visible'].append(is_visible)
            
            elif not is_visible and current_window is not None:
                # 结束当前窗口
                current_window.end_time = current_time
                # 设置窗口持续时间
                current_window.duration = current_time - current_window.start_time
                
                # 添加轨迹点数据
                if hasattr(observer, 'orbit') and current_window_points:
                    current_window.trajectory_points = current_window_points
                    current_window_points = None
                
                windows.append(current_window)
                current_window = None
            
            current_time += step
        
        # 处理最后一个未结束的窗口
        if current_window is not None:
            current_window.end_time = current_time
            # 设置窗口持续时间
            current_window.duration = current_time - current_window.start_time
            
            # 添加轨迹点数据
            if hasattr(observer, 'orbit') and current_window_points:
                current_window.trajectory_points = current_window_points
            
            windows.append(current_window)
        
        # 把全局轨迹点也保存到第一个窗口，用于整体可视化
        if windows and hasattr(observer, 'orbit'):
            # 复制第一个窗口
            first_window = windows[0]
            # 添加全局轨迹点
            first_window.trajectory_points['all_trajectory'] = all_trajectory_points
        
        return windows
    
    @staticmethod
    def calculate_constellation_access(constellation, targets, start_time, end_time, step=10*u.s, min_elevation=5.0):
        """
        计算星座和多个目标之间的访问窗口
        
        参数:
            constellation: 卫星星座对象
            targets: 目标列表或目标组对象
            start_time: 开始时间
            end_time: 结束时间
            step: 时间步长
            min_elevation: 最小仰角(度)，默认为5度
            
        返回:
            二维访问窗口列表，每个卫星对每个目标的访问窗口
        """
        # 如果targets是GroundTargetGroup，获取其targets列表
        if hasattr(targets, 'targets'):
            targets = targets.targets
            
        # 为每个卫星和每个目标计算访问窗口
        results = []
        
        for satellite in constellation.satellites:
            satellite_results = []
            
            for target in targets:
                # 计算卫星对地面目标的访问窗口
                windows = AccessCalculator.calculate_access_windows(
                    observer=satellite,  # 卫星作为观测者
                    observed=target,       # 地面目标作为观测目标
                    start_time=start_time,
                    end_time=end_time,
                    step=step,
                    min_elevation=min_elevation
                )
                satellite_results.append((target, windows))
                
            results.append((satellite, satellite_results))
            
        return results

    @staticmethod
    def _process_satellites_chunk(observers:List[Observer], observeds:List[Observed], start_time, end_time, step, min_elevation=5.0):
        """处理一组卫星的访问计算（作为静态方法支持多进程）"""
        chunk_results = []
        
        for observer in observers:
            satellite_results = []
            for observed in observeds:
                windows = AccessCalculator.calculate_access_windows(
                    observer=observer,
                    observed=observed,
                    start_time=start_time,
                    end_time=end_time,
                    step=step,
                    min_elevation=min_elevation
                )
                from satellite import Satellite
                from target import GroundTarget
                ground_instance = observed if isinstance(observed, GroundTarget) else observer
                satellite_results.append((ground_instance, windows))
            satellite_instance = observer if isinstance(observer, Satellite) else observed
            chunk_results.append((satellite_instance, satellite_results))
        
        return chunk_results
    
    @staticmethod
    def calculate_constellation_access_multiprocessing(
            constellation,
            targets,
            start_time,
            end_time,
            step=10*u.s,
            max_processes=None,
            min_elevation=5.0
        ):
        """使用multiprocessing模块计算星座和多个目标之间的访问窗口"""
        import multiprocessing as mp
        
        # 如果targets是GroundTargetGroup，获取其targets列表
        if hasattr(targets, 'targets'):
            targets = targets.targets
        
        # 预计算所有卫星的轨道数据
        precomputed_data = {}
        # 使用进程池 预计算所有卫星的轨道数据
        with mp.Pool(processes=len(constellation.satellites)) as pool:
            # 准备参数
            args = [(satellite, (start_time, end_time), step) for satellite in constellation.satellites]
            # 启动异步
            from satellite import Satellite
            results_async = pool.starmap_async(
                Satellite.precompute_satellite_orbits,
                args
            )
            # 等待并获取结果
            chunk_results = results_async.get()
            # 合并结果
            for chunk_result in chunk_results:
                satellite, obrit_data = chunk_result
                precomputed_data[satellite] = obrit_data
        # 将计算任务分配给多个进程
        num_processes = max_processes or mp.cpu_count()
        num_satellites = len(constellation.satellites)
        
        # 每个进程处理的卫星数量
        satellites_per_process = max(1, num_satellites // num_processes)
        
        # 分割卫星列表
        satellites = list(precomputed_data.keys())
        satellite_chunks = []
        for i in range(0, num_satellites, satellites_per_process):
            chunk = satellites[i:i+satellites_per_process]
            satellite_chunks.append(chunk)
        
        # 使用进程池
        with mp.Pool(processes=min(len(satellite_chunks), num_processes)) as pool:
            # 准备参数
            args = [(
                chunk,
                targets,
                start_time, 
                end_time,
                step,
                min_elevation,
            ) for chunk in satellite_chunks]
            
            # 启动异步计算
            results_async = pool.starmap_async(AccessCalculator._process_satellites_chunk, args)
            
            # 等待并获取结果
            chunk_results = results_async.get()
            
            # 打印进度
            print(f"计算完成: 100% ({len(satellite_chunks)}/{len(satellite_chunks)})")
        
        # 合并结果
        results = []
        for chunk_result in chunk_results:
            results.extend(chunk_result)

        return results